#include "build_kernels.h"
#include <utility\cuda_utils.h>
#include <device_launch_parameters.h>
#include <device_functions.h>

namespace gc {

	inline __device__ void atomicMerge(BV_SOA bv1, int idx1, const BV_SOA bv2, int idx2) {
		atomicMinCustom<Float>(bv1.minx + idx1, bv2.minx[idx2]);
		atomicMinCustom<Float>(bv1.miny + idx1, bv2.miny[idx2]);
		atomicMinCustom<Float>(bv1.minz + idx1, bv2.minz[idx2]);
		atomicMaxCustom<Float>(bv1.maxx + idx1, bv2.maxx[idx2]);
		atomicMaxCustom<Float>(bv1.maxy + idx1, bv2.maxy[idx2]);
		atomicMaxCustom<Float>(bv1.maxz + idx1, bv2.maxz[idx2]);
	}

	inline __device__ void fillBV(AABB_SOA& bv_soa, uint idx, const AABB& bv) {
		bv_soa.minx[idx] = bv.ptMin.x;
		bv_soa.miny[idx] = bv.ptMin.y;
		bv_soa.minz[idx] = bv.ptMin.z;
		bv_soa.maxx[idx] = bv.ptMax.x;
		bv_soa.maxy[idx] = bv.ptMax.y;
		bv_soa.maxz[idx] = bv.ptMax.z;
	}
	/*
	__forceinline__ __device__ unsigned dynamic_smem_size()
	{
		unsigned ret;
		asm volatile ("mov.u32 %0, %dynamic_smem_size;" : "=r"(ret));
		return ret;
	}*/

	__global__ void buildBV_SOA(int size, BV_SOA bvs, const SimpleSphere *prims) {
		int idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= size) return;

		const BV bv = prims[idx].calcBV();
		fillBV(bvs, idx, bv);
	}

	__global__ void buildBVs(int size, LBvhLeaves leaves, const SimpleSphere *prims) {
		int idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= size) return;

		leaves.primIdx[idx] = idx;
		BV bv = prims[idx].calcBV();
		fillBV(leaves.bv, idx, bv);
	}

	/// parallel reduce in O(logn)
	__global__ void calcSceneBV_1(int size, LBvhLeaves leaves, BV* d_out) {
		uint idx = threadIdx.x + blockIdx.x * blockDim.x;
		uint tid = threadIdx.x;
		extern __shared__ BV sdata[]; //use shared mem to avoid writing global mem
		if (idx < size) sdata[tid] = leaves.bv.getBV(idx);
		else sdata[tid] = BV();
		__syncthreads();

		for (uint s = blockDim.x / 2; s>0; s >>= 1) {
			if (tid < s) {
				sdata[tid].merge(sdata[tid + s]);
			}
			__syncthreads(); //ensure all merges at one iteration are done
		}
		if (tid == 0) {
			d_out[blockIdx.x] = sdata[tid];
		}
	}

	__global__ void calcSceneBV_2(int size, BV* d_in, BV* d_out) {
		uint idx = threadIdx.x + blockIdx.x * blockDim.x;
		uint tid = threadIdx.x;
		extern __shared__ BV sdata[]; //use shared mem to avoid writing global mem
		if (idx < size) sdata[tid] = d_in[idx];
		else sdata[tid] = BV();
		__syncthreads();

		for (uint s = blockDim.x / 2; s>0; s >>= 1) {
			if (tid < s) {
				sdata[tid].merge(sdata[tid + s]);
			}
			__syncthreads(); //ensure all merges at one iteration are done
		}
		if (tid == 0) {
			d_out[blockIdx.x] = sdata[tid];
		}
	}

	__global__ void calcMortonCodes_32b(int size, LBvhLeaves leaves, BV* scene) {
		uint idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= size) return;

		Point3f offset = (leaves.bv.center(idx) - scene->ptMin);
		offset.x /= (scene->ptMax.x - scene->ptMin.x);
		offset.y /= (scene->ptMax.y - scene->ptMin.y);
		offset.z /= (scene->ptMax.z - scene->ptMin.z);
		uint mc32 = morton3f_32b(offset.x, offset.y, offset.z);
		leaves.morton[idx] = mc32;
	}

	__global__ void reorderLeaves(int size, LBvhLeaves leaves, LBvhLeaves leaves_sorted, int *mapidx) {
		uint idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= size) return;

		int idx_old = mapidx[idx];
		leaves_sorted.primIdx[idx] = idx_old;
		BV bv = leaves.bv.getBV(idx_old);
		fillBV(leaves_sorted.bv, idx, bv);
		//leaves_sorted.bv[idx] = leaves.bv[idx_old];
		leaves_sorted.morton[idx] = leaves.morton[idx];
	}

	__global__ void calcBuildMetrics(int interiorSize, LBvhLeaves leaves, uint *metrics) {
		uint idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= interiorSize) return;

		// reflects the highest differing bit between the keys covered by interior node idx
		metrics[idx] = 32 - __clz(leaves.morton[idx] ^ leaves.morton[idx + 1]);
	}

	
	__global__ void buildInteriors(int leafSize, LBvhLeaves leaves, LBvhInteriors interiors, InteriorBuildAid aid) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= leafSize) return;

		int parent;
		if (idx == 0 || (idx != leafSize - 1 && aid.metrics[idx - 1] > aid.metrics[idx])) {
			// is left child of Interior idx
			parent = idx;
			leaves.parent[idx] = parent;
			interiors.lc[parent] = idx;
			interiors.rangeL[parent] = idx;
			//printf("interior %d has rangeL %d\n", parent, idx);
			leaves.segLen[idx] = 1;
			aid.leftLeafCount[idx] = 1;
			//atomicAdd(leaves.segLen + idx, 1);
			//atomicAdd(aid.leftLeafCount + idx, 1);
		}
		else {
			// is right child of Interior idx-1
			parent = idx - 1;
			leaves.parent[idx] = parent;
			interiors.rc[parent] = idx;
			interiors.rangeR[parent] = idx;
			//printf("interior %d has rangeR %d\n", parent, idx);
			//printf("buildInteriors: rangeR of %d = %d\n", parent, idx);
		}

		// choose parent in bottom-up fashion. O(n)
		int current, left, right;
		int _lc, _rc, lc, rc;
		bool isleaf_lc, isleaf_rc;
		BV_SOA bv_l, bv_r;
		while (atomicAdd(aid.visitCount + parent, 1) == 1) {
			current = parent;
			// merge bv
			_lc = interiors.lc[current], _rc = interiors.rc[current];
			isleaf_lc = isLeaf(_lc, &lc);
			isleaf_rc = isLeaf(_rc, &rc);
			if (!isleaf_lc && !isleaf_rc) {
				bv_l = interiors.bv;
				bv_r = interiors.bv;
			}
			else if (!isleaf_lc && isleaf_rc) {
				bv_l = interiors.bv;
				bv_r = leaves.bv;
			}
			else if (isleaf_lc && !isleaf_rc) {
				bv_l = leaves.bv;
				bv_r = interiors.bv;
			}
			else {
				bv_l = leaves.bv;
				bv_r = leaves.bv;
			}
			interiors.bv.minx[current] = fmin(bv_l.minx[lc], bv_r.minx[rc]);
			interiors.bv.miny[current] = fmin(bv_l.miny[lc], bv_r.miny[rc]);
			interiors.bv.minz[current] = fmin(bv_l.minz[lc], bv_r.minz[rc]);
			interiors.bv.maxx[current] = fmax(bv_l.maxx[lc], bv_r.maxx[rc]);
			interiors.bv.maxy[current] = fmax(bv_l.maxy[lc], bv_r.maxy[rc]);
			interiors.bv.maxz[current] = fmax(bv_l.maxz[lc], bv_r.maxz[rc]);

			left = interiors.rangeL[current];
			right = interiors.rangeR[current];
			//if (right==0) printf("rangeR of %d in buildInteriors = 0\n", current);
			__threadfence();

			if (left == 0 && right == leafSize - 1) {
				//root
				break;
			}

			if (left == 0 || (right != leafSize - 1 && aid.metrics[left - 1] > aid.metrics[right])) {
				// is left child of Interior right
				parent = right;
				interiors.parent[current] = parent;
				interiors.lc[parent] = interiorIdx(current);
				interiors.rangeL[parent] = left;
				//printf("interior %d has rangeL %d\n", parent, left);

				//aid.leftLeafCount[parent] = old;
				leaves.segLen[left] ++;  // add LCL value of the leftmost leaf
				//atomicAdd(leaves.segLen + left, 1);
				aid.leftLeafCount[parent] = leaves.segLen[left]; // count the order of Interior parent
				__threadfence();
			}
			else {
				// is right child of Interior left-1
				parent = left - 1;
				interiors.parent[current] = parent;
				interiors.rc[parent] = interiorIdx(current);
				interiors.rangeR[parent] = right;
				__threadfence();
			}
		}
	}
	

	/*
	__global__ void buildInteriors(int leafSize, LBvhLeaves leaves, LBvhInteriors interiors, InteriorBuildAid aid) {
		int idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= leafSize) return;

		int parent;
		if (idx == 0 || (idx != leafSize - 1 && aid.metrics[idx - 1] > aid.metrics[idx])) {
			// is left child of Interior idx
			parent = idx;
			leaves.parent[idx] = parent;
			interiors.lc[parent] = idx;
			interiors.rangeL[parent] = idx;
			leaves.segLen[idx] = 1;
			aid.leftLeafCount[idx] = 1;
			interiors.bv.setBV(parent, leaves.bv.getBV(idx));

			//aid.visitCount[parent] = 1;
			__stwb(aid.visitCount + parent, 1);
			__threadfence();
			//printf("left leaf %d visited parent %d\n", idx, parent);
			return;
		}

		// is right child of Interior idx-1
		parent = idx - 1;
		//printf("right leaf %d visited parent %d\n", idx, parent);
		leaves.parent[idx] = parent;
		interiors.rc[parent] = idx;
		interiors.rangeR[parent] = idx;
		interiors.bvR.setBV(parent, leaves.bv.getBV(idx));
		__threadfence();

		int flag;
		do {
			flag = __ldca(aid.visitCount + parent);
			//flag = aid.visitCount[parent];
		} while (flag == 0);

	    //printf("right leaf idx %d enters loop\n",idx);

		// choose parent in bottom-up fashion. O(n)
		int current, left, right;
		while (true) {
			current = parent;
			left = interiors.rangeL[current];
			right = interiors.rangeR[current];
			interiors.bv.merge(current, interiors.bvR.getBV(current));

			if (left == 0 && right == leafSize - 1) {
				//root
				//interiors.parent[current] = -1;
				break;
			}

			if (left == 0 || (right != leafSize - 1 && aid.metrics[left - 1] > aid.metrics[right])) {
				// is left child of Interior right
				parent = right;
				interiors.parent[current] = parent;
				interiors.lc[parent] = interiorIdx(current);
				interiors.rangeL[parent] = left;
				interiors.bv.setBV(parent, interiors.bv.getBV(current));

				leaves.segLen[left] ++;  // add LCL value of the leftmost leaf
				aid.leftLeafCount[parent] = leaves.segLen[left]; // count the order of Interior parent

				//aid.visitCount[parent] = 1;
				__stwb(aid.visitCount + parent, 1);
				__threadfence();
				break;
			}

			// is right child of Interior left-1
			parent = left - 1;
			interiors.parent[current] = parent;
			interiors.rc[parent] = interiorIdx(current);
			interiors.rangeR[parent] = right;
			interiors.bvR.setBV(parent, interiors.bv.getBV(current));
			__threadfence();
			int flag;
			do {
				//flag = aid.visitCount[parent];
				flag = __ldca(aid.visitCount + parent);
			} while (flag == 0);
		}
	}
	*/

	__global__ void calcLeafEscape(int size, LBvhLeaves leaves, LBvhInteriors interiors) {
		int idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= size) return;

		if (idx == size - 1) {
			leaves.escape[idx] = INT_MIN; // nothing to traverse next
			//printf("leaf %d has escape %d\n", idx, INT_MIN);
			return;
		}

		leaves.escape[idx] = interiors.rc[idx];
		//printf("leaf %d has escape %d\n", idx, interiors.rc[idx]);
	}

	__global__ void calcInteriorEscape(int size, LBvhLeaves leaves, LBvhInteriors interiors) {
		int idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= size) return;

		int rightmostLeaf = interiors.rangeR[idx];

		if (rightmostLeaf == size) {
			interiors.escape[idx] = INT_MIN;
			//printf("interior %d has escape %d\n", idx, INT_MIN);
			return;
		}
		interiors.escape[idx] = interiors.rc[rightmostLeaf];
		//printf("interior %d has escape %d\n", idx, interiors.rc[rightmostLeaf]);
	}

	__global__ void calcInteriorNewIdx(int size, LBvhLeaves leaves, LBvhInteriors interiors, int *leftLeafCount, int *mapidx) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		if (idx >= size) return;

		int binIdx = interiors.rangeL[idx];
		int idx_new = leaves.segOffset[binIdx] + leaves.segLen[binIdx] - leftLeafCount[idx];
		mapidx[idx] = idx_new;
	}

	__global__ void reorderInteriors(int interiorSize, LBvhLeaves leaves, LBvhInteriors interiors, LBvhInteriors interiors_sorted, int *mapidx) {
		int idx = blockIdx.x * blockDim.x + threadIdx.x;
		// process interior
		if (idx >= interiorSize) return;
		int idx_new = mapidx[idx];

		interiors_sorted.rangeL[idx_new] = interiors.rangeL[idx];
		interiors_sorted.rangeR[idx_new] = interiors.rangeR[idx];
		BV bv = interiors.bv.getBV(idx);
		fillBV(interiors_sorted.bv, idx_new, bv);
	}

	__global__ void resetBVs(int size, AABB_SOA bv_soa) {
		int idx = blockIdx.x*blockDim.x + threadIdx.x;
		if (idx >= size) return;
		bv_soa.maxx[idx] = -FMAX;
		bv_soa.maxy[idx] = -FMAX;
		bv_soa.maxz[idx] = -FMAX;
		bv_soa.minx[idx] = FMAX;
		bv_soa.miny[idx] = FMAX;
		bv_soa.minz[idx] = FMAX;
	}

}